import cv2
import numpy as np
from sklearn.cluster import KMeans
import os

pic_dir = '../img/other'
pic_name = '0614-a-rgb_color.bmp'

def extract_main_colors(image_path, num_colors=3):
    # 检查文件是否存在
    if not os.path.exists(image_path):
        raise FileNotFoundError(f"文件不存在：{image_path}")

    # 读取图片
    img = cv2.imread(image_path)

    # 检查图片是否正确加载
    if img is None:
        raise ValueError(f"无法加载图像：{image_path}")

    # 将图片转换为RGB格式
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    # 将图片转换为numpy数组并重塑为二维数组
    img_array = np.reshape(img, (img.shape[0] * img.shape[1], 3))

    # 使用K-means聚类提取主要颜色
    kmeans = KMeans(n_clusters=num_colors, random_state=0, n_init=10)  # 显式设置n_init参数
    kmeans.fit(img_array)

    # 获取聚类中心（即主要颜色）
    main_colors = kmeans.cluster_centers_

    # 将颜色转换为整数类型
    main_colors = main_colors.astype(np.int32)  # 确保颜色值是整数类型

    return main_colors

def display_main_colors(image_path, main_colors):
    # 创建一个空白图像用于显示主色调
    color_display = np.zeros((100, 300, 3), dtype=np.uint8)

    # 计算每个颜色块的宽度
    color_width = 300 // len(main_colors)

    # 在空白图像上绘制主色调
    for i, color in enumerate(main_colors):
        # 确保颜色值在0到255之间
        color = np.clip(color, 0, 255)
        color = tuple(color.tolist())  # 确保颜色是元组格